package meiyunting.shiwuzhaoling_gatway.filter;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义限流过滤器工厂
 */
@Slf4j
@Component
public class RateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<RateLimiterGatewayFilterFactory.Config> {

    private final Map<String, TokenBucket> buckets = new ConcurrentHashMap<>();

    public RateLimiterGatewayFilterFactory() {
        super(Config.class);
    }

    @Override
    public List<String> shortcutFieldOrder() {
        return Arrays.asList("replenishRate", "burstCapacity");
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            ServerHttpRequest request = exchange.getRequest();
            ServerHttpResponse response = exchange.getResponse();

            // 获取IP地址作为限流键
            String ip = request.getRemoteAddress().getHostString();
            
            // 获取或创建令牌桶
            TokenBucket bucket = buckets.computeIfAbsent(ip, 
                key -> new TokenBucket(config.getReplenishRate(), config.getBurstCapacity()));
            
            // 尝试获取令牌
            boolean allowed = bucket.tryConsume();
            
            if (allowed) {
                // 允许请求通过
                return chain.filter(exchange).then(Mono.fromRunnable(() -> {
                    // 添加限流相关的响应头
                    Map<String, String> headers = bucket.getHeaders();
                    headers.forEach((key, value) -> response.getHeaders().add(key, value));
                }));
            } else {
                // 请求被限流，返回429状态码
                log.warn("请求被限流: {} - {}", ip, request.getURI().getPath());
                response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                return response.setComplete();
            }
        };
    }

    /**
     * 配置类
     */
    @Data
    public static class Config {
        // 令牌补充速率（每秒）
        private int replenishRate = 10;
        // 令牌桶容量
        private int burstCapacity = 20;
    }

    /**
     * 令牌桶实现
     */
    private static class TokenBucket {
        private final int replenishRate;
        private final int burstCapacity;
        private final AtomicInteger tokens;
        private long lastRefillTime;
        
        public TokenBucket(int replenishRate, int burstCapacity) {
            this.replenishRate = replenishRate;
            this.burstCapacity = burstCapacity;
            this.tokens = new AtomicInteger(burstCapacity);
            this.lastRefillTime = Instant.now().toEpochMilli();
        }
        
        public synchronized boolean tryConsume() {
            refill();
            if (tokens.get() > 0) {
                tokens.decrementAndGet();
                return true;
            }
            return false;
        }
        
        private void refill() {
            long now = Instant.now().toEpochMilli();
            long timeSinceLastRefill = now - lastRefillTime;
            
            // 计算需要补充的令牌数
            int tokensToAdd = (int) (timeSinceLastRefill / 1000.0 * replenishRate);
            if (tokensToAdd > 0) {
                int newTokens = Math.min(tokens.get() + tokensToAdd, burstCapacity);
                tokens.set(newTokens);
                lastRefillTime = now;
            }
        }
        
        public Map<String, String> getHeaders() {
            Map<String, String> headers = new ConcurrentHashMap<>();
            headers.put("X-RateLimit-Remaining", String.valueOf(tokens.get()));
            headers.put("X-RateLimit-Burst-Capacity", String.valueOf(burstCapacity));
            headers.put("X-RateLimit-Replenish-Rate", String.valueOf(replenishRate));
            return headers;
        }
    }
} 